#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time         : 19-2-9 下午9:06
# @Author       : ding
# @File         : tfconv.py
# @Description  : 转化图像为tfrecords
import codecs
import os
import sys

import tensorflow as tf
from abc import ABC, abstractmethod

from PIL import Image

sys.path.append("./ibase")

from ibase.ithread import IThread, auto_list


def int64_feature(values):
    if not isinstance(values, (tuple, list)):
        values = [values]
    return tf.train.Feature(int64_list=tf.train.Int64List(value=values))


def bytes_feature(values):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[values]))


def float_feature(values):
    if not isinstance(values, (tuple, list)):
        values = [values]
    return tf.train.Feature(float_list=tf.train.FloatList(value=values))


class TFConv:
    def __init__(self):
        self.feature_function = {"float": float_feature,
                                 "int64": int64_feature,
                                 "bytes": bytes_feature}

    def create_example(self, data_dict):
        feature = {}
        for key in data_dict.keys():
            data_type = data_dict[key]["type"]
            data_vals = data_dict[key]["data"]
            feature[key] = self.feature_function[data_type](data_vals)
        return tf.train.Example(features=tf.train.Features(feature=feature))

    def process_example(self, keys, idx, lock, bar, logger, idict, pattern):
        with tf.python_io.TFRecordWriter(pattern % idx) as f:
            for key in keys:
                try:
                    example = self.create_example(idict[key])
                    f.write(example.SerializeToString())
                except:
                    logger.error("key error : %s" % key)
                IThread.upd_bar(lock, bar)

    def run(self, idict, nthread, tfroot="./", pattern="object_train_number_of_%3d.tfrecords"):
        keys = list(idict.keys())

        tfpath = os.path.join(tfroot, pattern)
        result = auto_list(keys, nthread, target=self.process_example, same_args=(idict, tfpath), name="tf_conv")

        return result


class TFDict(ABC):
    def __init__(self, mode='RGB'):
        self.mode = mode

    def to_image_sample(self, img_dir, idict={}):
        image = Image.open(img_dir)
        if image.mode != self.mode:
            image = image.convert(self.mode)
        img_raw = image.tobytes()

        idict["image/encode"] = {"data": img_raw, "type": "bytes"}
        idict["image/shape"] = {"data": [image.size[1], image.size[0], 3], "type": "int64"}

        return idict

    @staticmethod
    def to_key_val(val, key, tftype, idict={}):
        idict["image/%s" % key] = {"data": val, "type": tftype}

        return idict

    @abstractmethod
    def process(self, ilst, idx, lock, bar, logger):
        pass

        return None

    def list_to_dict(self, list_vals, nthread):
        idict = {}
        result = auto_list(list_vals, nthread, target=self.process, name="to_dict")

        for dval in result:
            try:
                idict.update(dval)
            except:
                pass
        return idict


class TFDictNorm(TFDict):
    def process(self, ilst, idx, lock, bar, logger):
        idict = {}
        for lst in ilst:
            try:
                img_dir, img_lab = lst.strip().split()
                img_dict = self.to_image_sample(img_dir)
                lab_dict = self.to_key_val(int(img_lab), "label", "int64")

                key = os.path.basename(img_dir)
                idict[key] = {}
                idict[key].update(img_dict)
                idict[key].update(lab_dict)
                IThread.upd_bar(lock, bar)
            except:
                logger.error(lst)
        return idict


if __name__ == "__main__":
    tfdict = TFDictNorm()

    with codecs.open("/home/ding/models/datas/disk/fine_data/OCR/ID_CARD/idcard_labeled_v1_867.txt", "r", encoding='utf-8') as f:
        lst_val = f.readlines()

    idict = tfdict.list_to_dict(lst_val, 4)
    # for key in idict.keys():
    #     print(idict[key])

    tfconv = TFConv()
    tfconv.run(idict, 4, pattern="idcard_train_001_of_%03d.tfrecords")

